#############################################################################################
#############################################################################################
# 
#   Open Source License/Disclaimer, Forecast Systems Laboratory NOAA/OAR/GSD, 
#   325 Broadway Boulder, CO 80305
#
#   This software is distributed under the Open Source Definition, which may be 
#   found at http://www.opensource.org/.
#   In particular, redistribution and use in source and binary forms, with or 
#   without modification, are permitted provided that the following conditions are met:
#
#   - Redistributions of source code must retain this notice, this list of 
#     conditions and the following disclaimer.
#
#   - Redistributions in binary form must provide access to this notice, this 
#     list of conditions and the following disclaimer, and the underlying source code.
#
#   - All modifications to this software must be clearly documented, and are 
#     solely the responsibility of the agent making the modifications.
#
#   - If significant modifications or enhancements are made to this software, 
#     the GSD Software Policy Manager (softwaremgr.fsl@noaa.gov) should be notified.
#
#   THIS SOFTWARE AND ITS DOCUMENTATION ARE IN THE PUBLIC DOMAIN AND ARE 
#   FURNISHED "AS IS." THE AUTHORS, THE UNITED STATES GOVERNMENT, ITS INSTRUMENTALITIES,
#   OFFICERS, EMPLOYEES, AND AGENTS MAKE NO WARRANTY, EXPRESS OR IMPLIED, AS TO
#   THE USEFULNESS OF THE SOFTWARE AND DOCUMENTATION FOR ANY PURPOSE. THEY ASSUME 
#   NO RESPONSIBILITY (1) FOR THE USE OF THE SOFTWARE AND DOCUMENTATION; OR (2) TO PROVIDE
#   TECHNICAL SUPPORT TO USERS. 
#############################################################################################
#############################################################################################

import os, copy

try:
    import numpy as np
except:
    raise ImportError('The numpy library cannot be found!')

try:
    from cdar.models.Shallow_Water_2D.ShallowWater2DModel import ShallowWaterModelClass as model
except:
    raise ImportError('The model library cannot be found!')
        
try:    
    from cdar.utilities import namelist_read as nml
except:
    raise ImportError('The nml library cannot be found!')

try:
    import ESMF
except:
    raise ImportError('The ESMP library cannot be found!')
        
EnSize = 100
SampleTime = 1800

DataFilePath='./data'  # Path to store or read FG, B and OBS

# Read in parameters for model configurations
params = nml.namelist_read(DataFilePath)

# start up ESMF
# this call is not necessary unless you want to to override the
# default options:
#  LogKind = NONE
#  debug = False
manager = ESMF.Manager(logkind=ESMF.LogKind.MULTI, debug=True)

# inquire for rank and proc from ESMF Virtual Machine
localPet = manager.local_pet
petCount = manager.pet_count
print "localPet = %d and petCount = %d " % (localPet, petCount)

# opening remarks
if localPet == 0:
    print "Welcome to the Ensemble GEN BE Utility!"

# Set up the model    
truth = model(params)  

# Print domain info 
if localPet == 0:
    print truth  
    
# Redefine the history output interval
truth.history_interval=SampleTime
truth.model_output_history()

# Check if SampleTime is in history
if not SampleTime in truth.history:
    raise Exception ("SampleTime is not in model history !")
    
#vis.plot_3d(truth.x, truth.y, truth.h, 'H at Time ='+str(0))

#output = truth.model_integration()
#u,v,h = np.split(output[Ntimestep-1,:],truth.VarDim)
#vis.plot_3d(truth.x, truth.y, h, 'H at Time ='+str(stop_time))

"""
from cdar.models.Lorenz_63.Lorenz63 import Lorenz63ModelClass as model
import cdar.models.Lorenz_63.Lorenz63Vis as vis


start_time = 0
stop_time = 4
Ntimestep = 1000

truth = model(start_time,stop_time,Ntimestep,time_diff_choice='leapfrog',x_noise=0.1,y_noise=0.1, \
                 z_noise=0.1,ic_start_range=-1,ic_end_range=1) 
                 
"""
#
#   No Change below
#                                  
traj = np.empty( (EnSize,truth.VarDim*(truth.nx),truth.ny) )

for i in range(EnSize):
    domain = copy.copy(truth)
    domain.model_ic_add_noise()
    domain.ic = domain.ic_perturbation
    # Array traj is the shape of [EnSize, VarDim*nx, ny]
    output = domain.model_integration()
    for j in range(len(domain.history)):
        if domain.history[j] == SampleTime:
            traj[i] =  domain.strip_off_boundaries( output[j] )

# Calculate mean
traj_mean = (1./float(EnSize))*np.tile(traj.sum(0), (EnSize,1,1))
#print 'traj_mean shape is', traj_mean.shape

# Calculate ensemble perturbation from mean for SampleTime
dtraj = traj - traj_mean
dtraj = dtraj[:,:,:].reshape((EnSize,dtraj[0,:,:].size))
#print 'dtraj shape is', dtraj.shape

# Calculate error covariance matrix
# COV is (VarDim)x(VarDim)
COV = (1./float(EnSize-1))*np.dot(dtraj.transpose(),dtraj)
di = np.diag_indices(dtraj.shape[1])
COV[di] = COV[di] + 0.0000001
#COV[di] = COV[di] + 1.0

# Compute inv(COV)*dD
# Should be (VarDim)x(VarDim)
Binv = np.linalg.solve( COV, np.eye(int(np.sqrt(COV.size))) )

print 'Is the Binv good ? ', np.allclose( np.dot(COV, Binv), np.eye(int(np.sqrt(Binv.size))) )   

if localPet == 0:
    print "The ensemble size is %d" % EnSize
    print "The Binv size is %d x %d" % ((int(np.sqrt(Binv.size)),)*2)
    print "Save the BE in file: %s" %  os.path.join(DataFilePath,'B_'+truth.Name+'_'+str(EnSize)+'_'+str(truth.nx)+'_'+str(SampleTime))+'.npz'
    np.savez(os.path.join(DataFilePath,'B_'+truth.Name+'_'+str(EnSize)+'_'+str(truth.nx)+'_'+str(SampleTime)), B=Binv)

    print 'B=\n',Binv
        
"""
# Pack the arraries on 3rd axis [Ntimestep, VarDim, EnSize]
y_t = traj[0,:,:,:].copy()

for i in range(1,EnSize):
    x_t = np.dstack((y_t,traj[i,:,:,:]))
    y_t = x_t.copy()

# Rearrange the array to [EnSize, Ntimestep, VarDim]
x_t = np.swapaxes(np.swapaxes(x_t,2,0),1,2)

vis.movie(x_t,N_trajectories=EnSize)
"""